/**
 * 
 */
package asm;

import java.io.*;
import java.util.NoSuchElementException;
import java.util.Scanner;
import java.util.TreeMap;
import java.util.ArrayList;

/**
 * The main assembler class.
 * Takes an input file name as a parameter. Creates an output file with the
 * same file name but a .out extension.
 * @author Piotr
 *
 */
public class Assembler {

	private static File sourceFile;
	private static File outputFile;
	private static PrintWriter outputWriter;
	private static TreeMap<String, Integer> symbolTable;
	private static ArrayList<Object[]> branchTable;
	private static TreeMap<String, String> mnemonicTable;
	private static TreeMap<String, Integer> regAliasTable;
	private static ArrayList<String> outputText;
	
	/**
	 * Main function. Takes input file name as parameter.
	 * @param args The input file name.
	 */
	public static void main(String[] args) throws IOException {

		initTables();
		// Try to open the source file
		if (args.length > 0) {
			sourceFile = new File(args[0]);
			if (!sourceFile.exists()) {
				System.out.println("File \"" + args[0] + "\" does not exist");
				System.exit(1);
			}
		createOutputFile(args[0]);
		translateFile();
		writeOutputFile();		//Throws IOException
			
		outputWriter.close();	//Throws IOException
		} else {
			System.out.println("No input file specified");
			System.exit(1);
		}
	}
	




	/**
	 * Writes the contents of the outputText array out to the output file.
	 * @throws IOException 
	 */
	private static void writeOutputFile() throws IOException {
		int lineNum = 0;
		for (String line: outputText ) {
			outputWriter.println(lineNum + "\t" + line);
			lineNum++;
		}
		outputWriter.flush();
	}






	/**
	 * Attempts to create an output file with the same name as the input file
	 * but with a .out extension. (stripping the original extension)
	 */
	private static void createOutputFile(String inputFileName) {
		String outputFileName;
		//Scan the input file name from the end, to identify the extension
		//Assume a blank extension if the beginning of the string is reached
		//or a pathname delimiter.
		int ii = inputFileName.length();
		char ch;
		do {
			ii--;
			ch = inputFileName.charAt(ii);
		} while (ii > 0 && ch != '.' && ch != '/' && ch != '\\');
		if (ii == 0 || ch == '/' || ch == '\\') {
			outputFileName = inputFileName + ".out";
		} else {
			outputFileName = inputFileName.substring(0, ii) + ".out";
		}
		//System.out.println(outputFileName);
		outputFile = new File(outputFileName);
		try {
			outputWriter = new PrintWriter(outputFile);
		} catch (IOException e) {
			System.out.println("Error creating output file \"" + outputFileName + "\"");
			System.exit(1);
		}
	}

	/**
	 * Subroutine that performs the overall function of translating the input file
	 * in assembly to an output file in machine code.
	 */
	private static void translateFile() throws IOException {
		symbolTable = new TreeMap<String, Integer>();
		branchTable = new ArrayList<Object[]>();
		outputText = new ArrayList<String>();
		//Read input file in line-by-line, create output file and branch/symbol tables
		doFirstPass();
		doSecondPass();
		
	}
	
	/**
	 * Method to perform the second pass of the assembler.
	 * Iterates through the branch table, calculates the branch targets using
	 * the symbol table, and fills in the target addresses.
	 */
	private static void doSecondPass() {
		for (Object[] entry: branchTable) {
			int source = (Integer)entry[0];
			String label = (String)entry[1];
			if (symbolTable.containsKey(label)) {
				int target = symbolTable.get(label);
				int offset = target - (source + 1);
				String offsetStr = toBits(offset, 8, true);
				if (offsetStr == null) {
					System.out.println("Error: Label \"" + label + 
					"\" out of range.");
				}
				String code = outputText.get(source);
				code = code.replace("#target#", offsetStr);
				outputText.set(source, code);
			} else {
				System.out.println("Error: Label \"" + label + 
				"\" not found.");
			}
			
			
		}
		
	}






	/**
	 * Subroutine to do the first pass of the assembler. 
	 * @return	The number of lines in the output file.
	 * @throws FileNotFoundException 
	 */
	private static int doFirstPass() throws FileNotFoundException {
		Scanner scanIn = new Scanner(sourceFile);
		String line;
		int inLineCount = 0;
		int outLineCount = 0;
		while (scanIn.hasNext()) {
			line = scanIn.nextLine();
			inLineCount++;
			outLineCount = doFirstPassLine(inLineCount, outLineCount, line);
		}
		return outLineCount;
	}
	
	/**
	 * Subroutine to do the first-pass processing on a line from the input file
	 * @param inLineCout The current line of the input file
	 * @param outLineCount The current line count in the output file
	 * @param line The String of the line from the input to be processed
	 * @return lineCount incremented by the number of lines that were added to the output
	 */
	private static int doFirstPassLine(int inLineCnt, int lineCount, String line) {
		Scanner scanLine = new Scanner(line);
		String token;
		if (!scanLine.hasNext()) {
			return lineCount;	//The string is empty
		} else {
			//Check the beginning of a line for a label
			//Any word followed by a colon found at the beginning of the line is
			//treated as a label.
			if ((token = scanLine.findInLine("^\\s*\\w+\\s*:")) != null) {
				token = token.split(":")[0];
				token = token.trim();
				//System.out.println("Found label: " + token);
				//Check if the label is a mnemonic
				if (mnemonicTable.containsKey(token.toLowerCase())) {
					System.out.println("Syntax error: Line " + inLineCnt + 
							" Invalid label.");
					System.exit(1);
				}
				//Add label to the symbol table
				if (!symbolTable.containsKey(token)) {
					symbolTable.put(token, lineCount);
				} else {
					System.out.println("Syntax error: Line " + inLineCnt + 
					" Invalid label reuse.");
					System.exit(1);
				}
			}
			//Check any following token for a mnemonic
			if (scanLine.hasNext()) {
				token = scanLine.next();
				//System.out.println(token);
				if (mnemonicTable.containsKey(token.toLowerCase())) {
					int jj;
					jj = doMnemonicTranslation(inLineCnt, lineCount, token.toLowerCase(),
							scanLine, line);
					return lineCount + jj;
				} else {
					//Check for comments, else print error
					if (token.startsWith(";") || token.startsWith("//")) {
						//it's a comment, ignore it
					} else {
						System.out.println("Syntax error: Line " + inLineCnt + 
						" Mnemonic, comment, or label expected.");
						System.exit(1);
					}
				}
			}
			
		}
		return lineCount;
	}

	/**
	 * Helper method to translate a single line, given a valid mnemonic and the scanner
	 * pointing after the mnemonic. Adds a line to the outputText, and fills in the
	 * branchTable if necessary.
	 * @param lineNum The line number of the output file (address)
	 * @param mnemonic
	 * @param scanLine Scanner pointing after the mnemonic
	 * @return returns the number of output lines created.
	 */
	private static int doMnemonicTranslation(int inLnCt, int lineNum, String mnemonic,
											Scanner scanLine, String line) {
		String mn = mnemonic;
		String rs;
		String rt;
		String rd;
		//int offset;
		//String target;
		//String comment = "";
		String outCode = "";
		int numOutLines = 0;
		//System.out.println(mn);
		try {
			if (mn.equals("add") || mn.equals("sub") || mn.equals("and") || mn.equals("or")
					|| mn.equals("nor") || mn.equals("mul") ) {	// All 3 operand instructions
				//System.out.println("Caught mnemonic");
				String[] temp = scanLine.nextLine().split(",|//|;",4);
				//System.out.println("temp.length " + temp.length);
				if (temp.length < 3) {
					throw new NoSuchElementException();
				}
				rd = temp[0].trim();
				rs = temp[1].trim();
				rt = temp[2].trim();
				//comment = (temp.length >= 4 ? temp[3] : "").trim();
				//System.out.println("rd " + rd);
				//System.out.println("rs " + rs);
				//System.out.println("rt " + rt);
				//System.out.println("rem " + comment);
				outCode = mnemonicTable.get(mn) 
				             + toBits(regAliasTable.get(rd), 4, false)
				             + toBits(regAliasTable.get(rs), 4, false)
				             + toBits(regAliasTable.get(rt), 4, false);
				//System.out.println(outCode);
				numOutLines = 1;
			} else if (mn.equals("sll") || mn.equals("srl") || mn.equals("sra")
					|| mn.equals("addi")) {
				String[] temp = scanLine.nextLine().split(",|//|;",4);
				//System.out.println("temp.length " + temp.length);
				if (temp.length < 3) {
					throw new NoSuchElementException();
				}
				rd = temp[0].trim();
				rs = temp[1].trim();
				rt = temp[2].trim();
				//comment = (temp.length >= 4 ? temp[3] : "").trim();
				//System.out.println("rd " + rd);
				//System.out.println("rs " + rs);
				//System.out.println("rt " + rt);
				//System.out.println("rem " + comment);
				if (mn.equals("addi")) {
					rt = toBits(new Integer(rt), 4, true);
				} else { // shifts
					rt = toBits((new Integer(rt)) - 1, 4, false);
				}
				if (rt == null) {
					System.out.println("Error: Line " + inLnCt + ", value out of range.");
				}
				outCode = mnemonicTable.get(mn) 
				             + toBits(regAliasTable.get(rd), 4, false)
				             + toBits(regAliasTable.get(rs), 4, false)
				             + rt;
				//System.out.println(outCode);
				numOutLines = 1;
			} else if (mn.equals("li")) {
				String[] temp = scanLine.nextLine().split(",|//|;",4);
				//System.out.println("temp.length " + temp.length);
				if (temp.length < 2) {
					throw new NoSuchElementException();
				}
				rd = temp[0].trim();
				rs = temp[1].trim();
				//comment = (temp.length >= 3 ? temp[2] : "").trim();
				//System.out.println("rd " + rd);
				//System.out.println("rs " + rs);
				//System.out.println("rt " + rt);
				//System.out.println("rem " + comment);
				rs = toBits(new Integer(rs), 8, false);
				if (rs == null) {
					System.out.println("Error: Line " + inLnCt + ", value out of range.");
				}
				outCode = mnemonicTable.get(mn) 
				             + toBits(regAliasTable.get(rd), 4, false)
				             + rs;
				//System.out.println(outCode);
				numOutLines = 1;
			} else if (mn.equals("load") || mn.equals("store")) {
				String[] temp = scanLine.nextLine().split(",",2);
				//System.out.println("temp.length " + temp.length);
				if (temp.length < 2) {
					throw new NoSuchElementException();
				}
				rd = temp[0].trim();
				temp = temp[1].split("\\(|\\)|//|;");
				if (temp.length < 2) {
					throw new NoSuchElementException();
				}
				rt = temp[0].trim();
				rs = temp[1].trim();
				//comment = (temp.length >= 4 ? temp[3] : "").trim();
				//System.out.println("rd " + rd);
				//System.out.println("rs " + rs);
				//System.out.println("rt " + rt);
				//System.out.println("rem " + comment);
				rt = toBits(new Integer(rt), 4, true);
				if (rt == null) {
					System.out.println("Error: Line " + inLnCt + ", value out of range.");
				}
				outCode = mnemonicTable.get(mn) 
				             + toBits(regAliasTable.get(rd), 4, false)
				             + toBits(regAliasTable.get(rs), 4, false)
				             + rt;
				//System.out.println(outCode);
				numOutLines = 1;
			} else if (mn.equals("bez") || mn.equals("blz")) {
				String[] temp = scanLine.nextLine().split(",|//|;",4);
				//System.out.println("temp.length " + temp.length);
				if (temp.length < 2) {
					throw new NoSuchElementException();
				}
				rd = temp[0].trim();
				rs = temp[1].trim();
				//comment = (temp.length >= 3 ? temp[2] : "").trim();
				//System.out.println("rd " + rd);
				//System.out.println("rs " + rs);
				//System.out.println("rt " + rt);
				//System.out.println("rem " + comment);
				outCode = mnemonicTable.get(mn) 
				             + toBits(regAliasTable.get(rd), 4, false)
				             + "#target#";
				Object[] bTableE = {new Integer(lineNum), rs};
				branchTable.add(bTableE);
				//System.out.println(outCode);
				numOutLines = 1;
			} else { // mn must be jalr
				String[] temp = scanLine.nextLine().split(",|//|;",4);
				//System.out.println("temp.length " + temp.length);
				if (temp.length < 2) {
					throw new NoSuchElementException();
				}
				rd = temp[0].trim();
				rs = temp[1].trim();
				//comment = (temp.length >= 3 ? temp[2] : "").trim();
				//System.out.println("rd " + rd);
				//System.out.println("rs " + rs);
				//System.out.println("rt " + rt);
				//System.out.println("rem " + comment);
				outCode = mnemonicTable.get(mn) 
				             + toBits(regAliasTable.get(rd), 4, false)
				             + toBits(regAliasTable.get(rs), 4, false)
				             + "0000";
				//System.out.println(outCode);
				numOutLines = 1;
			}
			outputText.add(outCode + "\t" + line.trim() ); 
		} catch (NoSuchElementException e) {
			System.out.println("Syntax error: Line " + inLnCt + ", operand expected.");
			System.exit(1);
		} catch (NullPointerException e) {
			System.out.println("Syntax error: Line " + inLnCt + ", operand error.");
			System.exit(1);
		}
		return numOutLines;
		
	}
	
	
	/**
	 * Helper method. Creates a string of bits given an integer and the length.
	 * Either sign-extends the string.
	 * @param integer The number to be converted to bits
	 * @param length The length of the string to be created
	 * @param signed If true, the number will be sign-extended
	 * @return Returns a null if it cannot be represented in the specified length
	 */
	private static String toBits(Integer integer, int length, boolean signed) {
		String out = Integer.toBinaryString(integer);
		if (signed) {
			if (integer < 0) {
				if (integer < (-1 * (int)Math.pow(2, (length - 1)))) {
					return null;
				} else {
					out = out.substring(out.length() - length);
				}
			} else {	//Positive signed
				if (integer > (int)Math.pow(2,(length - 1)) - 1) {
					return null;
				} else {
					out = charSeq('0', length-out.length()) + out;
				}
			}
		} else { //Unsigned
			if (integer < 0 || integer > (int)Math.pow(2,(length)) - 1) {
				return null;
			} else {
				out = charSeq('0', length-out.length()) + out; 
			}
		}
		return out;
	}
	
	/**
	 * Returns a string composed of ch repeated num times
	 * @param ch
	 * @param num
	 */
	private static String charSeq(char ch, int num) {
		String out = "";
		for (int ii = 0 ; ii < num ; ii++) {
			out = out + ch;
		}
		return out;
	}

	
	/**
	 * Fills in the tables for the mnemonics and register aliases
	 */
	private static void initTables() {
		mnemonicTable = new TreeMap<String, String>();
		regAliasTable = new TreeMap<String, Integer>();
		
		// Add supported instructions here. First field is opcode, second field is funct
		mnemonicTable.put("add",	"0000");
		mnemonicTable.put("sub",	"0001");
		mnemonicTable.put("and",	"0010");
		mnemonicTable.put("or", 	"0011");
		mnemonicTable.put("nor",	"0100");
		mnemonicTable.put("mul",	"0101");
		mnemonicTable.put("sll",	"0110");
		mnemonicTable.put("srl",	"0111");
		mnemonicTable.put("sra",	"1000");
		mnemonicTable.put("li", 	"1001");
		mnemonicTable.put("addi",	"1010");
		mnemonicTable.put("load",	"1011");
		mnemonicTable.put("store",	"1100");
		mnemonicTable.put("bez",	"1101");
		mnemonicTable.put("blz",	"1110");
		mnemonicTable.put("jalr",	"1111");
		mnemonicTable.put("nop",	"0000");
		
		// All aliases for registers
		regAliasTable.put("r0", 0);
		regAliasTable.put("r1", 1);
		regAliasTable.put("r2", 2);
		regAliasTable.put("r3", 3);
		regAliasTable.put("r4", 4);
		regAliasTable.put("r5", 5);
		regAliasTable.put("r6", 6);
		regAliasTable.put("r7", 7);
		regAliasTable.put("r8", 8);
		regAliasTable.put("r9", 9);
		regAliasTable.put("r10", 10);
		regAliasTable.put("r11", 11);
		regAliasTable.put("r12", 12);
		regAliasTable.put("r13", 13);
		regAliasTable.put("r14", 14);
		regAliasTable.put("r15", 15);
		regAliasTable.put("r16", 16);
		regAliasTable.put("r17", 17);
		regAliasTable.put("r18", 18);
		regAliasTable.put("r19", 19);
		regAliasTable.put("r20", 20);
		regAliasTable.put("r21", 21);
		regAliasTable.put("r22", 22);
		regAliasTable.put("r23", 23);
		regAliasTable.put("r24", 24);
		regAliasTable.put("r25", 25);
		regAliasTable.put("r26", 26);
		regAliasTable.put("r27", 27);
		regAliasTable.put("r28", 28);
		regAliasTable.put("r29", 29);
		regAliasTable.put("r30", 30);
		regAliasTable.put("r31", 31);
		
	}

	
	
}
